#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>
#include <fcntl.h>
#include <unistd.h>
#include <sys/stat.h>
#include <sys/uio.h>
#include <errno.h>
#include <arpa/inet.h>  // 字节序转换（htons/ntohs等）
#include "skiplist.h"

#define LOG(_fmt, ...) fprintf(stdout, "[%s:%d]: "_fmt, __FILE__, __LINE__, ##__VA_ARGS__)

// 原有随机层数生成函数不变
int random_level() {
    int level = 1;
    while ((rand() % 100) < P * 100 && level < MAX_LEVEL) {
        level++;
    }
    return level;
}

// 原有节点创建函数不变
SNode *create_sk_node(const char *key, const char *value, int level) {
    SNode *node = (SNode *)kvs_malloc(sizeof(SNode) + level * sizeof(SNode *));
    node->key = strdup(key);
    node->value = strdup(value);
    return node;
}

// 1. 修改初始化函数，添加锁初始化
SkipList *skiplist_init() {
    srand((unsigned int)time(NULL));
    SkipList *sl = (SkipList *)kvs_malloc(sizeof(SkipList));
    sl->level = 1;
    sl->size = 0;
    sl->header = (SNode *)kvs_malloc(sizeof(SNode) + MAX_LEVEL * sizeof(SNode *));
    sl->header->key = NULL;
    sl->header->value = NULL;
    for (int i = 0; i < MAX_LEVEL; i++) {
        sl->header->forward[i] = NULL;
    }

    // 初始化互斥锁（递归锁）
#if SKIPLIST_MUTEX_LOCK
    pthread_mutexattr_t mutex_attr;
    pthread_mutexattr_init(&mutex_attr);
    pthread_mutexattr_settype(&mutex_attr, PTHREAD_MUTEX_RECURSIVE);
    pthread_mutex_init(&sl->mutex_lock, &mutex_attr);
    pthread_mutexattr_destroy(&mutex_attr);
#endif

    // 初始化读写锁（写者优先）
#if SKIPLIST_RW_LOCK
    pthread_rwlockattr_t rw_attr;
    pthread_rwlockattr_init(&rw_attr);
#ifdef PTHREAD_RWLOCK_PREFER_WRITER_NONRECURSIVE_NP
    pthread_rwlockattr_setkind_np(&rw_attr, PTHREAD_RWLOCK_PREFER_WRITER_NONRECURSIVE_NP);
#elif defined(PTHREAD_RWLOCK_PREFER_WRITER)
    pthread_rwlockattr_setkind(&rw_attr, PTHREAD_RWLOCK_PREFER_WRITER);
#endif
    pthread_rwlock_init(&sl->rw_lock, &rw_attr);
    pthread_rwlockattr_destroy(&rw_attr);
#endif

    return sl;
}

// 2. 为原有操作函数添加锁（仅展示修改部分，核心逻辑不变）
int skiplist_insert(SkipList *sl, const char *key, const char *value) {
    if (!sl || !key || !value) return -1;

    // 加写锁
#if SKIPLIST_MUTEX_LOCK
    pthread_mutex_lock(&sl->mutex_lock);
#elif SKIPLIST_RW_LOCK
    pthread_rwlock_wrlock(&sl->rw_lock);
#endif

    SNode *update[MAX_LEVEL];
    SNode *current = sl->header;

    // 原有查找前驱节点逻辑不变
    for (int i = sl->level - 1; i >= 0; i--) {
        while (current->forward[i] != NULL && 
               strcmp(current->forward[i]->key, key) < 0) {
            current = current->forward[i];
        }
        update[i] = current;
    }

    current = current->forward[0];
    if (current != NULL && strcmp(current->key, key) == 0) {
        // 解锁后返回（键已存在）
#if SKIPLIST_MUTEX_LOCK
        pthread_mutex_unlock(&sl->mutex_lock);
#elif SKIPLIST_RW_LOCK
        pthread_rwlock_unlock(&sl->rw_lock);
#endif
        return 1;
    }

    int new_level = random_level();
    if (new_level > sl->level) {
        for (int i = sl->level; i < new_level; i++) {
            update[i] = sl->header;
        }
        sl->level = new_level;
    }

    SNode *new_node = create_sk_node(key, value, new_level);
    for (int i = 0; i < new_level; i++) {
        new_node->forward[i] = update[i]->forward[i];
        update[i]->forward[i] = new_node;
    }

    sl->size++;

    // 解锁
#if SKIPLIST_MUTEX_LOCK
    pthread_mutex_unlock(&sl->mutex_lock);
#elif SKIPLIST_RW_LOCK
    pthread_rwlock_unlock(&sl->rw_lock);
#endif

    return 0;
}

SNode *skiplist_search(SkipList *sl, const char *key) {
    if (!sl || !key) return NULL;

    // 加读锁
#if SKIPLIST_MUTEX_LOCK
    pthread_mutex_lock(&sl->mutex_lock);
#elif SKIPLIST_RW_LOCK
    pthread_rwlock_rdlock(&sl->rw_lock);
#endif

    SNode *current = sl->header;
    // 原有查找逻辑不变
    for (int i = sl->level - 1; i >= 0; i--) {
        while (current->forward[i] != NULL && 
               strcmp(current->forward[i]->key, key) < 0) {
            current = current->forward[i];
        }
    }

    current = current->forward[0];
    SNode *result = NULL;
    if (current != NULL && strcmp(current->key, key) == 0) {
        result = current;  // 找到节点
    }

    // 解锁
#if SKIPLIST_MUTEX_LOCK
    pthread_mutex_unlock(&sl->mutex_lock);
#elif SKIPLIST_RW_LOCK
    pthread_rwlock_unlock(&sl->rw_lock);
#endif

    return result;
}

int skiplist_update(SkipList *sl, const char *key, const char *new_value) {
    if (!sl || !key || !new_value) return -1;

    // 加写锁
#if SKIPLIST_MUTEX_LOCK
    pthread_mutex_lock(&sl->mutex_lock);
#elif SKIPLIST_RW_LOCK
    pthread_rwlock_wrlock(&sl->rw_lock);
#endif

    // 调用无锁查找（避免重复加锁）
    SNode *current = sl->header;
    for (int i = sl->level - 1; i >= 0; i--) {
        while (current->forward[i] != NULL && 
               strcmp(current->forward[i]->key, key) < 0) {
            current = current->forward[i];
        }
    }
    current = current->forward[0];

    int ret = 1;  // 默认未找到
    if (current != NULL && strcmp(current->key, key) == 0) {
        kvs_free(current->value);
        current->value = strdup(new_value);
        ret = 0;  // 修改成功
    }

    // 解锁
#if SKIPLIST_MUTEX_LOCK
    pthread_mutex_unlock(&sl->mutex_lock);
#elif SKIPLIST_RW_LOCK
    pthread_rwlock_unlock(&sl->rw_lock);
#endif

    return ret;
}

int skiplist_delete(SkipList *sl, const char *key) {
    if (!sl || !key) return -1;

    // 加写锁
#if SKIPLIST_MUTEX_LOCK
    pthread_mutex_lock(&sl->mutex_lock);
#elif SKIPLIST_RW_LOCK
    pthread_rwlock_wrlock(&sl->rw_lock);
#endif

    SNode *update[MAX_LEVEL];
    SNode *current = sl->header;

    // 原有查找前驱节点逻辑不变
    for (int i = sl->level - 1; i >= 0; i--) {
        while (current->forward[i] != NULL && 
               strcmp(current->forward[i]->key, key) < 0) {
            current = current->forward[i];
        }
        update[i] = current;
    }

    current = current->forward[0];
    if (current == NULL || strcmp(current->key, key) != 0) {
        // 解锁后返回（节点不存在）
#if SKIPLIST_MUTEX_LOCK
        pthread_mutex_unlock(&sl->mutex_lock);
#elif SKIPLIST_RW_LOCK
        pthread_rwlock_unlock(&sl->rw_lock);
#endif
        return 1;
    }

    // 原有删除节点逻辑不变
    for (int i = 0; i < sl->level; i++) {
        if (update[i]->forward[i] != current) break;
        update[i]->forward[i] = current->forward[i];
    }

    kvs_free(current->key);
    kvs_free(current->value);
    kvs_free(current);

    while (sl->level > 1 && sl->header->forward[sl->level - 1] == NULL) {
        sl->level--;
    }

    sl->size--;

    // 解锁
#if SKIPLIST_MUTEX_LOCK
    pthread_mutex_unlock(&sl->mutex_lock);
#elif SKIPLIST_RW_LOCK
    pthread_rwlock_unlock(&sl->rw_lock);
#endif

    return 0;
}

// 3. 修改获取大小函数，添加锁
int skiplist_get_size(SkipList *sl) {
    if (!sl) return 0;

    // 加读锁
#if SKIPLIST_MUTEX_LOCK
    pthread_mutex_lock(&sl->mutex_lock);
#elif SKIPLIST_RW_LOCK
    pthread_rwlock_rdlock(&sl->rw_lock);
#endif

    int size = sl->size;

    // 解锁
#if SKIPLIST_MUTEX_LOCK
    pthread_mutex_unlock(&sl->mutex_lock);
#elif SKIPLIST_RW_LOCK
    pthread_rwlock_unlock(&sl->rw_lock);
#endif

    return size;
}

// 4. 新增：清空跳表（保留结构，删除所有节点）
int skiplist_clear(SkipList *sl) {
    if (!sl) return -1;

    // 加写锁
#if SKIPLIST_MUTEX_LOCK
    pthread_mutex_lock(&sl->mutex_lock);
#elif SKIPLIST_RW_LOCK
    pthread_rwlock_wrlock(&sl->rw_lock);
#endif

    // 遍历底层节点，释放所有数据节点
    SNode *current = sl->header->forward[0];
    SNode *next;
    while (current != NULL) {
        next = current->forward[0];
        kvs_free(current->key);
        kvs_free(current->value);
        kvs_free(current);
        current = next;
    }

    // 重置跳表状态（表头指针、层数、大小）
    for (int i = 0; i < sl->level; i++) {
        sl->header->forward[i] = NULL;
    }
    sl->level = 1;
    sl->size = 0;

    // 解锁
#if SKIPLIST_MUTEX_LOCK
    pthread_mutex_unlock(&sl->mutex_lock);
#elif SKIPLIST_RW_LOCK
    pthread_rwlock_unlock(&sl->rw_lock);
#endif

    return 0;
}

// 5. 修改销毁函数，添加锁销毁
void skiplist_destroy(SkipList *sl) {
    if (!sl) return;

    // 加写锁（确保销毁时无并发操作）
#if SKIPLIST_MUTEX_LOCK
    pthread_mutex_lock(&sl->mutex_lock);
#elif SKIPLIST_RW_LOCK
    pthread_rwlock_wrlock(&sl->rw_lock);
#endif

    // 原有销毁逻辑不变
    SNode *current = sl->header->forward[0];
    SNode *next;
    while (current != NULL) {
        next = current->forward[0];
        kvs_free(current->key);
        kvs_free(current->value);
        kvs_free(current);
        current = next;
    }
    kvs_free(sl->header);

    // 销毁锁
#if SKIPLIST_MUTEX_LOCK
    pthread_mutex_unlock(&sl->mutex_lock);
    pthread_mutex_destroy(&sl->mutex_lock);
#elif SKIPLIST_RW_LOCK
    pthread_rwlock_unlock(&sl->rw_lock);
    pthread_rwlock_destroy(&sl->rw_lock);
#endif

    kvs_free(sl);
}

// ====================全量持久化相关函数====================
// 辅助函数：遍历跳表（底层有序）并写入文件
static int skiplist_traverse_persist(SkipList *sl, int fd) {
    if (!sl || fd == -1) return -1;

    // 遍历底层节点（跳表底层为有序链表）
    SNode *current = sl->header->forward[0];
    while (current != NULL) {
        int key_len = strlen(current->key);
        int val_len = strlen(current->value);

#if 1 // MASTER_SLAVE 为了方便各种模式都能使用正常同一个文件，这里统一存储/读取方式
        // 网络字节序转换
        uint32_t net_key_len = htonl(key_len);
        uint32_t net_val_len = htonl(val_len);
        struct iovec iov[4] = {
            {&net_key_len, 4},
            {&net_val_len, 4},
            {current->key, key_len},
            {current->value, val_len}
        };
#else
        struct iovec iov[4] = {
            {&key_len, 4},
            {&val_len, 4},
            {current->key, key_len},
            {current->value, val_len}
        };
#endif

        ssize_t total = writev(fd, iov, 4);
        if (total != 4 + 4 + key_len + val_len) {
            perror("skiplist traverse persist failed");
            return -1;
        }
        current = current->forward[0];
    }
    return 0;
}

// 全量持久化
int skiplist_full_store(SkipList *sl, const char *full_filename) {
    if (!sl || !full_filename) return -1;

    // 加读锁（仅读取数据）
#if SKIPLIST_MUTEX_LOCK
    pthread_mutex_lock(&sl->mutex_lock);
#elif SKIPLIST_RW_LOCK
    pthread_rwlock_rdlock(&sl->rw_lock);
#endif

    int fd = open(full_filename, O_WRONLY | O_CREAT | O_TRUNC, 0644);
    if (fd == -1) {
        perror("full store open failed");
#if SKIPLIST_MUTEX_LOCK
        pthread_mutex_unlock(&sl->mutex_lock);
#elif SKIPLIST_RW_LOCK
        pthread_rwlock_unlock(&sl->rw_lock);
#endif
        return -1;
    }

    int ret = skiplist_traverse_persist(sl, fd);
    fsync(fd);
    close(fd);

    // 解锁
#if SKIPLIST_MUTEX_LOCK
    pthread_mutex_unlock(&sl->mutex_lock);
#elif SKIPLIST_RW_LOCK
    pthread_rwlock_unlock(&sl->rw_lock);
#endif

    return ret;
}

// 带备份的全量持久化
int skiplist_full_store_with_backup(SkipList *sl, const char *full_filename) {
    if (!sl || !full_filename) return -1;

    // 加读锁
#if SKIPLIST_MUTEX_LOCK
    pthread_mutex_lock(&sl->mutex_lock);
#elif SKIPLIST_RW_LOCK
    pthread_rwlock_rdlock(&sl->rw_lock);
#endif

    // 备份旧文件
    struct stat st;
    if (stat(full_filename, &st) == 0) {
        char backup_filename[256];
        time_t now = time(NULL);
        struct tm* tm = localtime(&now);
        snprintf(backup_filename, sizeof(backup_filename),
                 "%s.%04d%02d%02d%02d%02d",
                 full_filename, tm->tm_year + 1900, tm->tm_mon + 1, tm->tm_mday,
                 tm->tm_hour, tm->tm_min);
        rename(full_filename, backup_filename);
    }

    // 写入新数据
    int fd = open(full_filename, O_WRONLY | O_CREAT | O_TRUNC, 0644);
    if (fd == -1) {
        perror("full store with backup open failed");
#if SKIPLIST_MUTEX_LOCK
        pthread_mutex_unlock(&sl->mutex_lock);
#elif SKIPLIST_RW_LOCK
        pthread_rwlock_unlock(&sl->rw_lock);
#endif
        return -1;
    }

    int ret = skiplist_traverse_persist(sl, fd);
    fsync(fd);
    close(fd);

    // 解锁
#if SKIPLIST_MUTEX_LOCK
    pthread_mutex_unlock(&sl->mutex_lock);
#elif SKIPLIST_RW_LOCK
    pthread_rwlock_unlock(&sl->rw_lock);
#endif

    return ret;
}

// 从文件恢复跳表
int skiplist_full_restore(const char *full_filename, SkipList *sl) {
    if (!full_filename || !sl) return -1;

    // 先清空现有数据
    if (skiplist_clear(sl) != 0) {
        fprintf(stderr, "clear before restore failed\n");
        return -1;
    }

    int fd = open(full_filename, O_RDONLY);
    if (fd == -1) {
        perror("restore open failed");
        return -1;
    }

#if 1 // MASTER_SLAVE 为了方便各种模式都能使用正常同一个文件，这里统一存储/读取方式
    uint32_t net_key_len, net_val_len;
#endif
    int key_len, val_len;
    char *key_buf = NULL, *val_buf = NULL;
    long int read_count = 0;
    int is_eof = 0;

    while (1) {
#if 1 // MASTER_SLAVE 为了方便各种模式都能使用正常同一个文件，这里统一存储/读取方式
        // 读网络字节序长度
        ssize_t n = read(fd, &net_key_len, 4);
        if (n == 0) { is_eof = 1; break; }
        if (n != 4) break;
        key_len = ntohl(net_key_len);

        n = read(fd, &net_val_len, 4);
        if (n == 0) { is_eof = 1; break; }
        if (n != 4) break;
        val_len = ntohl(net_val_len);
#else
        // 读主机字节序长度
        ssize_t n = read(fd, &key_len, 4);
        if (n == 0) { is_eof = 1; break; }
        if (n != 4) break;

        n = read(fd, &val_len, 4);
        if (n == 0) { is_eof = 1; break; }
        if (n != 4) break;
#endif

        // 校验长度
        if (key_len <= 0 || val_len <= 0) {
            fprintf(stderr, "invalid key/val length\n");
            break;
        }

        // 动态分配缓冲区
        key_buf = kvs_malloc(key_len + 1);
        if (!key_buf) { perror("kvs_malloc key_buf failed"); break; }

        val_buf = kvs_malloc(val_len + 1);
        if (!val_buf) {
            perror("kvs_malloc val_buf failed");
            kvs_free(key_buf);
            key_buf = NULL;
            break;
        }

        // 读键值
        n = read(fd, key_buf, key_len);
        if (n == 0) { is_eof = 1; break; }
        if (n != key_len) break;
        key_buf[key_len] = '\0';

        n = read(fd, val_buf, val_len);
        if (n == 0) { is_eof = 1; break; }
        if (n != val_len) break;
        val_buf[val_len] = '\0';

        // 插入跳表（内部已加锁）
        if (skiplist_insert(sl, key_buf, val_buf) == 0) {
            read_count++;
        }

        // 释放缓冲区
        kvs_free(key_buf);
        kvs_free(val_buf);
        key_buf = NULL;
        val_buf = NULL;
    }

    LOG("成功读取了%ld个节点\n", read_count);

    // 处理异常
    if (!is_eof && errno != 0 && errno != EINTR) {
        perror("restore read failed");
        skiplist_clear(sl); // 恢复失败则清空
    }

    // 清理资源
    if (key_buf) kvs_free(key_buf);
    if (val_buf) kvs_free(val_buf);
    close(fd);

    return 0;
}